from cozepy import COZE_CN_BASE_URL, Coze, TokenAuth, Stream, WorkflowEvent, WorkflowEventType

def run_coze_workflow(user_input: str, workflow_id: str, api_token: str, base_url: str = COZE_CN_BASE_URL):
    """
    运行Coze工作流并处理返回的流事件

    参数:
        user_input: 用户输入的内容
        workflow_id: 工作流ID
        api_token: Coze API访问令牌
        base_url: Coze API基础URL(默认为中国区)

    返回:
        str: 收集到的消息内容或错误信息
    """
    # 初始化Coze客户端
    coze = Coze(auth=TokenAuth(token=api_token), base_url=base_url)
    result = ""  # 用于存储最终结果

    def handle_workflow_iterator(stream: Stream[WorkflowEvent]):
        """处理工作流事件迭代器"""
        nonlocal result
        for event in stream:
            if event.event == WorkflowEventType.MESSAGE:
                message = event.message.content if event.message else ""
                result += message + "\n"  # 将消息追加到结果中
            elif event.event == WorkflowEventType.ERROR:
                error = event.error if event.error else "未知错误"
                result += "发生错误：" + str(error) + "\n"
            elif event.event == WorkflowEventType.INTERRUPT:
                handle_workflow_iterator(
                    coze.workflows.runs.resume(
                        workflow_id=workflow_id,
                        event_id=event.interrupt.interrupt_data.event_id,
                        resume_data="hey",
                        interrupt_type=event.interrupt.interrupt_data.type,
                    )
                )

    # 获取用户输入（当前由参数传入）

    # 运行工作流并处理结果
    handle_workflow_iterator(
        coze.workflows.runs.stream(
            workflow_id=workflow_id,
            parameters={"input": user_input}
        )
    )

    return result.strip()  # 去除首尾空白后返回结果


# 使用示例
if __name__ == "__main__":
    # 配置参数
    WORKFLOW_ID = "7516524156826746915"
    API_TOKEN = "pat_nnTFo3UOS0uEcaIOn4aY7GPHg1GjnFQgvX6WHtefmQW2jhnk7z6nt2tq2wbGFFws"

    # 运行工作流
    a=run_coze_workflow("数学",workflow_id=WORKFLOW_ID, api_token=API_TOKEN)
    print(a)